# predict_test.py

import os
# 关闭 tensorflow oneDNN 日志
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"   # 只保留 FATAL 日志

import matplotlib
matplotlib.use('TkAgg')
from matplotlib import rcParams
rcParams['font.sans-serif'] = ['SimHei']  # 中文字体
rcParams['axes.unicode_minus'] = False    # 负号正常显示

import pandas as pd
import numpy as np
import joblib
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import (
    classification_report, accuracy_score, f1_score,
    precision_score, recall_score, confusion_matrix,
    roc_auc_score, roc_curve
)
from word2vec_processor import Word2VecProcessor
from utils.config import Config


# ================= 配置路径 =================
root_path = "D:/pycode/group4_nlp_project/"
conf = Config(root_path)

data_path = conf.test_path   # 测试集路径
cat_model_path = conf.basic_model_path + "cat_rf.pkl"
cat_processor_path = conf.basic_model_path + "cat_word2vec.pkl"
label_model_path = conf.basic_model_path + "label_lr.pkl"
label_processor_path = conf.basic_model_path + "label_word2vec.pkl"


def evaluate_cat_model(df):
    print("\n================= 多分类任务 (cat) =================")
    # 加载模型和词向量
    rf = joblib.load(cat_model_path)
    processor = Word2VecProcessor.load(cat_processor_path)

    # 转换向量
    df["vector"] = df["review_clean"].apply(lambda x: processor.get_sentence_vector(str(x)))
    df = df[df["vector"].notnull()].reset_index(drop=True)
    X = np.vstack(df["vector"].values)
    y_true = df["cat"]

    # 预测
    y_pred = rf.predict(X)

    # 输出报告
    print(classification_report(y_true, y_pred))
    acc = accuracy_score(y_true, y_pred)
    f1_macro = f1_score(y_true, y_pred, average="macro")
    prec_macro = precision_score(y_true, y_pred, average="macro")
    rec_macro = recall_score(y_true, y_pred, average="macro")

    print(f"准确率: {acc:.4f}")
    print(f"宏平均F1: {f1_macro:.4f}")
    print(f"宏平均Precision: {prec_macro:.4f}")
    print(f"宏平均Recall: {rec_macro:.4f}")

    # 混淆矩阵
    cm = confusion_matrix(y_true, y_pred, labels=rf.classes_)
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
                xticklabels=rf.classes_, yticklabels=rf.classes_)
    plt.xlabel("预测值")
    plt.ylabel("真实值")
    plt.title("混淆矩阵 (cat - 随机森林)")
    plt.tight_layout()
    plt.show()


def evaluate_label_model(df):
    print("\n================= 二分类任务 (label) =================")
    # 加载模型和词向量
    lr = joblib.load(label_model_path)
    processor = Word2VecProcessor.load(label_processor_path)

    # 转换向量
    df["vector"] = df["review_clean"].apply(lambda x: processor.get_sentence_vector(str(x)))
    df = df[df["vector"].notnull()].reset_index(drop=True)
    X = np.vstack(df["vector"].values)
    y_true = df["label"]

    # 预测
    y_pred = lr.predict(X)
    y_prob = lr.predict_proba(X)[:, 1]

    # 输出报告
    print(classification_report(y_true, y_pred))
    acc = accuracy_score(y_true, y_pred)
    f1 = f1_score(y_true, y_pred)
    prec = precision_score(y_true, y_pred)
    rec = recall_score(y_true, y_pred)
    auc = roc_auc_score(y_true, y_prob)

    print(f"准确率: {acc:.4f}")
    print(f"F1: {f1:.4f}")
    print(f"Precision: {prec:.4f}")
    print(f"Recall: {rec:.4f}")
    print(f"AUC: {auc:.4f}")

    # 混淆矩阵
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(6, 5))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Greens",
                xticklabels=lr.classes_, yticklabels=lr.classes_)
    plt.xlabel("预测值")
    plt.ylabel("真实值")
    plt.title("混淆矩阵 (label - 逻辑回归)")
    plt.tight_layout()
    plt.show()

    # ROC曲线
    fpr, tpr, _ = roc_curve(y_true, y_prob)
    plt.figure(figsize=(6, 5))
    plt.plot(fpr, tpr, label=f"AUC={auc:.3f}")
    plt.plot([0, 1], [0, 1], "--", color="gray")
    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    plt.title("ROC曲线 (label - 逻辑回归)")
    plt.legend()
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    # 读测试集
    df_test = pd.read_csv(data_path)

    # 分别评估两个模型
    evaluate_cat_model(df_test)
    evaluate_label_model(df_test)
